from flask import Blueprint, jsonify, request
import os
import json
import pandas as pd
from datetime import datetime
import numpy as np

# 创建蓝图对象
strategy_combination = Blueprint('strategy_combination', __name__)

# 策略组合文件路径
STRATEGY_COMBINATION_FILE = 'data/strategy_combinations.json'

def ensure_data_file():
    """确保数据文件存在，如果不存在则创建一个空的数据文件"""
    os.makedirs(os.path.dirname(STRATEGY_COMBINATION_FILE), exist_ok=True)
    if not os.path.exists(STRATEGY_COMBINATION_FILE):
        with open(STRATEGY_COMBINATION_FILE, 'w', encoding='utf-8') as f:
            json.dump([], f, ensure_ascii=False)

def read_strategy_combinations():
    """读取所有策略组合"""
    ensure_data_file()
    with open(STRATEGY_COMBINATION_FILE, 'r', encoding='utf-8') as f:
        return json.load(f)

def write_strategy_combinations(data):
    """写入策略组合到文件"""
    ensure_data_file()
    with open(STRATEGY_COMBINATION_FILE, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)

@strategy_combination.route('/strategy_combinations', methods=['GET'])
def get_all_strategy_combinations():
    """获取所有策略组合"""
    try:
        combinations = read_strategy_combinations()
        return jsonify(combinations)
    except Exception as e:
        return jsonify({'error': f'获取策略组合失败: {str(e)}'}), 500

@strategy_combination.route('/strategy_combinations/<int:combination_id>', methods=['GET'])
def get_strategy_combination(combination_id):
    """获取单个策略组合"""
    try:
        combinations = read_strategy_combinations()
        if 0 <= combination_id < len(combinations):
            return jsonify(combinations[combination_id])
        else:
            return jsonify({'error': f'未找到组合ID: {combination_id}'}), 404
    except Exception as e:
        return jsonify({'error': f'获取策略组合失败: {str(e)}'}), 500

@strategy_combination.route('/strategy_combinations', methods=['POST'])
def add_strategy_combination():
    """添加新的策略组合"""
    try:
        new_combination = request.json
        if not new_combination or 'name' not in new_combination or 'strategies' not in new_combination:
            return jsonify({'error': '组合名称和策略列表是必需的参数'}), 400
        
        combinations = read_strategy_combinations()
        
        # 检查是否已存在相同名称的组合
        if any(c['name'] == new_combination['name'] for c in combinations):
            return jsonify({'error': f'组合名称已存在: {new_combination["name"]}'}), 409
        
        # 确保strategies是一个列表
        if not isinstance(new_combination['strategies'], list):
            return jsonify({'error': '策略列表必须是一个数组'}), 400
        
        # 添加创建时间
        new_combination['created_at'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        new_combination['updated_at'] = new_combination['created_at']
        
        combinations.append(new_combination)
        write_strategy_combinations(combinations)
        
        return jsonify({'id': len(combinations) - 1, **new_combination}), 201
    except Exception as e:
        return jsonify({'error': f'添加策略组合失败: {str(e)}'}), 500

@strategy_combination.route('/strategy_combinations/<int:combination_id>', methods=['PUT'])
def update_strategy_combination(combination_id):
    """更新策略组合"""
    try:
        updated_combination = request.json
        if not updated_combination or 'name' not in updated_combination or 'strategies' not in updated_combination:
            return jsonify({'error': '组合名称和策略列表是必需的参数'}), 400
        
        combinations = read_strategy_combinations()
        
        # 检查组合ID是否有效
        if not (0 <= combination_id < len(combinations)):
            return jsonify({'error': f'未找到组合ID: {combination_id}'}), 404
        
        # 检查是否与其他组合名称冲突
        if any(i != combination_id and c['name'] == updated_combination['name'] for i, c in enumerate(combinations)):
            return jsonify({'error': f'组合名称已存在: {updated_combination["name"]}'}), 409
        
        # 确保strategies是一个列表
        if not isinstance(updated_combination['strategies'], list):
            return jsonify({'error': '策略列表必须是一个数组'}), 400
        
        # 保留创建时间，更新修改时间
        updated_combination['created_at'] = combinations[combination_id].get('created_at', datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
        updated_combination['updated_at'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        
        combinations[combination_id] = updated_combination
        write_strategy_combinations(combinations)
        
        return jsonify({'id': combination_id, **updated_combination})
    except Exception as e:
        return jsonify({'error': f'更新策略组合失败: {str(e)}'}), 500

@strategy_combination.route('/strategy_combinations/<int:combination_id>', methods=['DELETE'])
def delete_strategy_combination(combination_id):
    """删除策略组合"""
    try:
        combinations = read_strategy_combinations()
        
        # 检查组合ID是否有效
        if not (0 <= combination_id < len(combinations)):
            return jsonify({'error': f'未找到组合ID: {combination_id}'}), 404
        
        deleted_combination = combinations.pop(combination_id)
        write_strategy_combinations(combinations)
        
        return jsonify({'message': f'已删除组合: {deleted_combination["name"]}', 'deleted': deleted_combination})
    except Exception as e:
        return jsonify({'error': f'删除策略组合失败: {str(e)}'}), 500

@strategy_combination.route('/strategy_combinations/<int:combination_id>/holdings', methods=['GET'])
def get_combination_holdings(combination_id):
    """获取策略组合的当前持仓情况"""
    try:
        from api.back_trade import TradingStrategy

        combinations = read_strategy_combinations()
        
        if not (0 <= combination_id < len(combinations)):
            return jsonify({'error': f'未找到组合ID: {combination_id}'}), 404
        
        # 获取查询参数
        start_date = request.args.get('start_date')
        is_real_param = request.args.get('isReal')  # 新增：从查询参数获取isReal
        
        combination = combinations[combination_id]
        strategies = combination['strategies']
        
        holdings = []
        total_positions_map = {}
        
        for strategy in strategies:
            parameters = strategy.get('parameters', {})
            
            # 确定是否使用实盘数据，查询参数优先，然后是策略参数
            if is_real_param is not None:
                is_real_for_strategy = is_real_param.lower() == 'true'
            else:
                is_real_for_strategy = bool(parameters.get('isReal', False))
            
            # 使用与回测相同的参数处理逻辑
            processed_params = {
                'ischeckhs300': bool(parameters.get('ischeckhs300', False)),
                'hs300_limit': int(parameters.get('hs300_limit', 30)),
                'backCycle': str(parameters.get('backCycle', 'day')),
                'codes': parameters.get('codes', []),
                'increaseday_range': int(parameters.get('increaseday_range', 21)),
                'highday_range': int(parameters.get('highday_range', 25)),
                'lowday_range': int(parameters.get('lowday_range', 17)),
                'sortType': int(parameters.get('sortType', 1)),
                'tradeType': int(parameters.get('tradeType', 1)),
                'isReal': is_real_for_strategy
            }
            
            # 如果提供了开始时间，添加到参数中
            if start_date:
                processed_params['startDate'] = start_date
            
            # 运行回测来获取最新持仓
            try:
                strategy_instance = TradingStrategy(processed_params)
                trade_log, cumulative_returns = strategy_instance.execute()
                
                # 直接从策略实例获取最终持仓
                positions = []
                if strategy_instance.position and len(strategy_instance.position) > 0:
                    positions = [strategy_instance.position]
                
                # 计算最终收益率
                final_return = 0
                if not cumulative_returns.empty and len(cumulative_returns) > 0:
                    last_return = cumulative_returns.iloc[-1]['cumulative_return']
                    # 检查是否为NaN
                    if pd.isna(last_return):
                        print(f"警告：策略 {strategy.get('name')} 的累计收益为NaN，设置为0")
                        final_return = 0
                    else:
                        final_return = (last_return - 1) * 100
                
                strategy_holdings = {
                    'strategy_name': strategy.get('name', '未命名策略'),
                    'parameters': parameters,
                    'positions': positions,
                    'return': final_return
                }
                holdings.append(strategy_holdings)
                
                # 合并到总持仓
                for pos in positions:
                    total_positions_map[pos['code']] = pos

            except Exception as e:
                print(f"获取策略 '{strategy.get('name')}' 持仓失败: {str(e)}")
                holdings.append({
                    'strategy_name': strategy.get('name', '未命名策略'),
                    'parameters': parameters,
                    'positions': [],
                    'return': 0,
                    'error': f'获取持仓失败: {str(e)}'
                })

        # 计算总收益（简单平均）
        valid_returns = [h.get('return', 0) for h in holdings if h.get('return') is not None and not (isinstance(h.get('return'), float) and np.isnan(h.get('return')))]
        total_return = sum(valid_returns) / len(valid_returns) if valid_returns else 0
        
        # 计算总持仓的占比
        total_positions_with_weight = []
        total_strategies = len(holdings)
        
        # 统计每个股票被多少个策略持有
        stock_strategy_count = {}
        for strategy_holding in holdings:
            for position in strategy_holding.get('positions', []):
                code = position['code']
                if code not in stock_strategy_count:
                    stock_strategy_count[code] = 0
                stock_strategy_count[code] += 1
        
        # 为总持仓添加占比信息
        for code, position in total_positions_map.items():
            strategy_count = stock_strategy_count.get(code, 0)
            weight = (strategy_count / total_strategies) * 100  # 转换为百分比
            
            position_with_weight = {
                **position,
                'weight': weight,
                'strategy_count': strategy_count
            }
            total_positions_with_weight.append(position_with_weight)
        
        combination_holdings = {
            'name': combination['name'],
            'created_at': combination.get('created_at', ''),
            'updated_at': combination.get('updated_at', ''),
            'strategies_holdings': holdings,
            'total_positions': total_positions_with_weight,
            'total_return': total_return,
            'start_date': start_date,  # 添加开始时间到返回结果中
            'is_real': is_real_param.lower() == 'true' if is_real_param else None  # 添加数据源信息
        }
        
        return jsonify(combination_holdings)
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        return jsonify({'error': f'获取组合持仓失败: {str(e)}'}), 500

@strategy_combination.route('/strategy_combinations/<int:combination_id>/backtest', methods=['POST'])
def run_combination_backtest(combination_id):
    """运行策略组合的回测"""
    try:
        from api.back_trade import get_data, TradingStrategy
        
        combinations = read_strategy_combinations()
        
        # 检查组合ID是否有效
        if not (0 <= combination_id < len(combinations)):
            return jsonify({'error': f'未找到组合ID: {combination_id}'}), 404
        
        combination = combinations[combination_id]
        strategies = combination['strategies']
        
        # 获取回测参数
        backtest_params = request.get_json(silent=True) or {}
        print(f"收到的回测参数: {backtest_params}")
        
        start_date = backtest_params.get('start_date')
        end_date = backtest_params.get('end_date')
        is_real_override = backtest_params.get('is_real')
        
        print(f"解析后: start_date={start_date}, end_date={end_date}, is_real_override={is_real_override}")
        
        # 运行每个策略的回测
        results = []
        combined_returns_data = []
        
        print(f"--- Strategies for combination {combination_id}: {strategies} ---") # 调试输出
        
        for strategy in strategies:
            print(f"--- Processing strategy: {strategy} ---") # 调试输出
            # 获取策略参数
            parameters = strategy.get('parameters', {})
            print(f"--- Strategy parameters: {parameters}, type: {type(parameters)} ---") # 调试输出
            
            # 确定是否使用实盘数据，UI的全局设置优先
            if is_real_override is not None:
                is_real_for_strategy = is_real_override
            else:
                is_real_for_strategy = bool(parameters.get('isReal', False))

            # 确保参数格式正确
            processed_params = {
                'ischeckhs300': bool(parameters.get('ischeckhs300', False)),
                'hs300_limit': int(parameters.get('hs300_limit', 30)),
                'backCycle': str(parameters.get('backCycle', 'day')),
                'codes': parameters.get('codes', []),
                'increaseday_range': int(parameters.get('increaseday_range', 21)),
                'highday_range': int(parameters.get('highday_range', 25)),
                'lowday_range': int(parameters.get('lowday_range', 17)),
                'sortType': int(parameters.get('sortType', 1)),
                'tradeType': int(parameters.get('tradeType', 1)),
                'isReal': is_real_for_strategy
            }
            
            # 参数验证：确保窗口大小不为0
            if processed_params['lowday_range'] <= 0:
                processed_params['lowday_range'] = 1
                print(f"警告：策略 {strategy.get('name')} 的lowday_range不能为0或负数，已设置为1")
            if processed_params['highday_range'] <= 0:
                processed_params['highday_range'] = 1
                print(f"警告：策略 {strategy.get('name')} 的highday_range不能为0或负数，已设置为1")
            
            # 添加开始日期（如果有）
            if start_date:
                processed_params['startDate'] = start_date
            elif 'startDate' in parameters:
                processed_params['startDate'] = parameters['startDate']
                
            # 运行回测
            try:
                # 创建策略实例
                strategy_instance = TradingStrategy(processed_params)
                
                # 执行回测
                trade_log, cumulative_returns  = strategy_instance.execute()
                print(f"--- Trade log: {trade_log} ---") # 调试输出
                print(f"--- Cumulative returns: {cumulative_returns} ---") # 调试输出
                # 提取结果
                # trade_log = backtest_result.get('trade_log', [])
                # cumulative_returns = backtest_result.get('cumulative_returns', [])
                
                # 计算最终收益率
                final_return = 0
                print(f"--- Cumulative returns length: {len(cumulative_returns)} ---") # 调试输出
                if not cumulative_returns.empty and len(cumulative_returns) > 0:
                    last_return = cumulative_returns.iloc[-1]['cumulative_return']
                    # 检查是否为NaN
                    if pd.isna(last_return):
                        print(f"警告：策略 {strategy.get('name')} 的累计收益为NaN，设置为0")
                        final_return = 0
                    else:
                        final_return = (last_return - 1) * 100
                print(f"--- Final return: {final_return} ---") # 调试输出
                strategy_result = {
                    'strategy_name': strategy.get('name', '未命名策略'),
                    'parameters': processed_params,
                    'trade_log': trade_log,
                    'cumulative_returns': cumulative_returns.to_dict('records') if not cumulative_returns.empty else [],
                    'final_return': final_return
                }
                print(f"--- Strategy result: {strategy_result} ---") # 调试输出
                results.append(strategy_result)
                
                # 添加到组合收益数据
                if not cumulative_returns.empty:
                    combined_returns_data.append(cumulative_returns)
            except Exception as e:
                print(f"运行策略 {strategy.get('name')} 的回测失败: {str(e)}")
                # 添加错误信息但继续处理其他策略
                results.append({
                    'strategy_name': strategy.get('name', '未命名策略'),
                    'parameters': processed_params,
                    'error': str(e),
                    'trade_log': [],
                    'cumulative_returns': [],
                    'final_return': 0
                })
        
        # 计算组合的累计收益（简单平均）
        combined_returns = []
        if combined_returns_data:
            # 创建日期到收益率的映射
            date_to_returns = {}
            
            # 收集所有日期和对应的收益率
            for strategy_returns in combined_returns_data:
                for _, row in strategy_returns.iterrows():
                    date = row['date']
                    if date not in date_to_returns:
                        date_to_returns[date] = []
                    date_to_returns[date].append(row['cumulative_return'])
            
            # 计算每个日期的平均收益率
            for date in sorted(date_to_returns.keys()):
                returns = date_to_returns[date]
                avg_return = sum(returns) / len(returns)
                combined_returns.append({
                    'date': date,
                    'cumulative_return': avg_return
                })
        
        return jsonify({
            'combination_name': combination['name'],
            'strategy_results': results,
            'combined_returns': combined_returns
        })
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        return jsonify({'error': f'运行组合回测失败: {str(e)}'}), 500 